package org.zjx.core;

import lombok.Data;
import lombok.SneakyThrows;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Logger;

/**
 * Socket连接池实现
 * 核心特性：
 * 1. 连接复用：减少TCP握手开销
 * 2. 空闲管理：自动回收超时连接
 * 3. 流量控制：限制最大并发连接数
 * 4. 连接验证：确保连接有效性
 */
public class SocketPool {

    // 连接池核心队列（线程安全）
    private final LinkedBlockingQueue<PooledSocket> pool;


    // 配置参数
    private final int maxTotal;       // 最大连接数
    private final long maxIdleTime;   // 最大空闲时间（毫秒）

    // 活跃连接计数器（原子操作）
    private final AtomicInteger activeCount = new AtomicInteger(0);

    /**
     * 连接池构造函数
     * @param maxTotal         最大连接数
     * @param maxIdleTime      最大空闲时间
     * @param unit             时间单位
     */
    public SocketPool(
                      int maxTotal,
                      long maxIdleTime,
                      TimeUnit unit) {
        this.maxTotal = maxTotal;
        this.maxIdleTime = unit.toMillis(maxIdleTime);
        this.pool = new LinkedBlockingQueue<>(maxTotal);
    }

    /**
     * 从连接池获取连接（带超时机制）
     * @return 可用的Socket连接
     */
    @SneakyThrows
    public PooledSocket borrowSocket(InetSocketAddress target){
        // 1. 尝试从池中获取空闲连接（最多等待5秒）
        PooledSocket socket = pool.poll(5, TimeUnit.SECONDS);

        // 2. 连接耗尽时创建新连接（不超过最大限制）
        if (socket == null) {
            if (activeCount.get() >= maxTotal) {
                throw new IOException("连接池已耗尽，当前最大连接数：" + maxTotal);
            }
            socket = createNewSocket(target);
        }

        // 3. 验证连接有效性
        if (!socket.isValid(maxIdleTime)) {
            closeSocket(socket);
            return borrowSocket(target); // 递归重试
        }

        return socket;
    }

    /**
     * 归还连接到连接池
     * @param socket 要归还的连接
     */
    public void returnSocket(PooledSocket socket) {
        if (socket.isValid(maxIdleTime)) {
            pool.offer(socket);  // 添加到池中
        } else {
            closeSocket(socket); // 无效连接直接关闭
        }
    }

    /**
     * 创建新Socket连接（带超时控制）
     */
    @SneakyThrows
    private PooledSocket createNewSocket(InetSocketAddress target){
        Socket s = new Socket();
        s.connect(target, 5000);    // 5秒连接超时
        s.setSoTimeout(10000);      // 10秒读取超时
        activeCount.incrementAndGet();
        return new PooledSocket(s);
    }

    /**
     * 安全关闭Socket连接
     */
    private void closeSocket(PooledSocket socket) {
        try {
            if (!socket.getSocket().isClosed()) {
                socket.getSocket().close();
            }
        } catch (IOException e) {
            // 记录关闭异常日志
        } finally {
            activeCount.decrementAndGet();
        }
    }

    /**
     * 清理所有连接（服务关闭时调用）
     */
    public void cleanup() {
        pool.forEach(this::closeSocket);
        pool.clear();
    }

    /**
     * 包装类：跟踪Socket状态和最后使用时间
     */
    @Data
    public static class PooledSocket {
        private final Socket socket;
        private long lastUsedTime;  // 最后使用时间戳

        public PooledSocket(Socket socket) {
            this.socket = socket;
            this.lastUsedTime = System.currentTimeMillis();
        }

        /**
         * 验证连接有效性
         * @return 是否有效
         */
        public boolean isValid(long maxIdleTime) {
            return !socket.isClosed() &&
                    (System.currentTimeMillis() - lastUsedTime) < maxIdleTime;
        }

        /**
         * 更新最后使用时间（调用时触发）
         */
        public void touch() {
            this.lastUsedTime = System.currentTimeMillis();
        }
    }
}